#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Auto Capture - Automated Hook and API Call
自动捕获加密过程
"""

import frida
import sys
import time
import threading
import requests
import json

captured_data = {
    'keys': [],
    'ivs': [],
    'encryptions': []
}

print("\n" + "="*80)
print("AUTO CAPTURE - Shumei Encryption Analysis")
print("="*80 + "\n")

# ============================================================================
# Frida Hook Setup
# ============================================================================

print("[1] Connecting to device...")
device = frida.get_usb_device()
print(f"    Connected: {device.name}")

print("[2] Attaching to PID 2586...")
session = device.attach(2586)
print("    Attached!")

print("[3] Loading hook script...")

hook_code = """
setTimeout(function() {
    Java.perform(function() {
        console.log("[HOOK] Ready");
    
    // Hook SecretKeySpec
    var SecretKeySpec = Java.use("javax.crypto.spec.SecretKeySpec");
    SecretKeySpec.$init.overload('[B', 'java.lang.String').implementation = function(key, algorithm) {
        var hexKey = "";
        for (var i = 0; i < key.length; i++) {
            hexKey += ("0" + (key[i] & 0xFF).toString(16)).slice(-2);
        }
        
        var strKey = "";
        for (var i = 0; i < key.length; i++) {
            var c = key[i] & 0xFF;
            strKey += (c >= 32 && c <= 126) ? String.fromCharCode(c) : ".";
        }
        
        send({
            type: 'key',
            algorithm: algorithm,
            length: key.length,
            hex: hexKey,
            str: strKey
        });
        
        return this.$init(key, algorithm);
    };
    
    // Hook IvParameterSpec
    var IvParameterSpec = Java.use("javax.crypto.spec.IvParameterSpec");
    IvParameterSpec.$init.overload('[B').implementation = function(iv) {
        var hexIv = "";
        for (var i = 0; i < iv.length; i++) {
            hexIv += ("0" + (iv[i] & 0xFF).toString(16)).slice(-2);
        }
        
        var strIv = "";
        for (var i = 0; i < iv.length; i++) {
            var c = iv[i] & 0xFF;
            strIv += (c >= 32 && c <= 126) ? String.fromCharCode(c) : ".";
        }
        
        send({
            type: 'iv',
            length: iv.length,
            hex: hexIv,
            str: strIv
        });
        
        return this.$init(iv);
    };
    
    // Hook Cipher.doFinal
    var Cipher = Java.use("javax.crypto.Cipher");
    Cipher.doFinal.overload('[B').implementation = function(input) {
        var mode = this.getAlgorithm();
        
        var hexInput = "";
        var maxLen = Math.min(64, input.length);
        for (var i = 0; i < maxLen; i++) {
            hexInput += ("0" + (input[i] & 0xFF).toString(16)).slice(-2);
        }
        
        var result = this.doFinal(input);
        
        var hexOutput = "";
        maxLen = Math.min(64, result.length);
        for (var i = 0; i < maxLen; i++) {
            hexOutput += ("0" + (result[i] & 0xFF).toString(16)).slice(-2);
        }
        
        send({
            type: 'cipher',
            algorithm: mode,
            inputLength: input.length,
            outputLength: result.length,
            inputHex: hexInput,
            outputHex: hexOutput
        });
        
        return result;
    };
    });
}, 0);
"""

script = session.create_script(hook_code)

def on_message(message, data):
    global captured_data
    
    if message['type'] == 'send':
        payload = message['payload']
        
        if isinstance(payload, str):
            print(f"    {payload}")
        elif isinstance(payload, dict):
            msg_type = payload.get('type')
            
            if msg_type == 'key':
                print(f"\n{'='*80}")
                print(f"[KEY CAPTURED]")
                print(f"{'='*80}")
                print(f"Algorithm: {payload['algorithm']}")
                print(f"Length: {payload['length']}")
                print(f"HEX: {payload['hex']}")
                print(f"STR: {payload['str']}")
                print(f"{'='*80}\n")
                captured_data['keys'].append(payload)
                
            elif msg_type == 'iv':
                print(f"\n[IV CAPTURED]")
                print(f"Length: {payload['length']}")
                print(f"HEX: {payload['hex']}")
                print(f"STR: {payload['str']}\n")
                captured_data['ivs'].append(payload)
                
            elif msg_type == 'cipher':
                print(f"\n[CIPHER OPERATION]")
                print(f"Algorithm: {payload['algorithm']}")
                print(f"Input Length: {payload['inputLength']}")
                print(f"Output Length: {payload['outputLength']}")
                print(f"Input (HEX): {payload['inputHex'][:64]}...")
                print(f"Output (HEX): {payload['outputHex'][:64]}...\n")
                captured_data['encryptions'].append(payload)
    
    elif message['type'] == 'error':
        print(f"\n[FRIDA ERROR] {message.get('description', 'Unknown error')}\n")

script.on('message', on_message)
script.load()
print("    Loaded!\n")

# ============================================================================
# API Call Function
# ============================================================================

def call_api():
    print("\n" + "="*80)
    print("[API] Calling send-test endpoint...")
    print("="*80 + "\n")
    
    try:
        response = requests.get("http://120.27.155.222:9999/api/message/send-test", timeout=10)
        
        if response.status_code == 200:
            data = response.json()
            print(f"[API Response]")
            print(f"  Code: {data.get('code')}")
            print(f"  Message: {data.get('msg')}")
            print(f"  Points: {data.get('points')}")
            
            if 'dev' in data:
                dev_value = data['dev']
                print(f"\n[ENCRYPTED DEV FIELD]")
                print(f"  Length: {len(dev_value)}")
                print(f"  Value: {dev_value}")
                
                # Decode base64 to see encrypted bytes
                import base64
                try:
                    encrypted_bytes = base64.b64decode(dev_value)
                    hex_encrypted = encrypted_bytes.hex()
                    print(f"  Hex (first 64): {hex_encrypted[:64]}...")
                    print(f"  Total bytes: {len(encrypted_bytes)}")
                except Exception as e:
                    print(f"  Decode error: {e}")
            
            print()
        else:
            print(f"[API Error] Status: {response.status_code}")
            
    except Exception as e:
        print(f"[API Error] {e}\n")

# ============================================================================
# Main Execution
# ============================================================================

print("[4] Waiting 3 seconds for hooks to settle...")
time.sleep(3)

print("\n[5] Making API calls...\n")

# Call API 3 times
for i in range(3):
    print(f"\n{'#'*80}")
    print(f"# API Call {i+1}/3")
    print(f"{'#'*80}")
    
    call_api()
    
    if i < 2:
        print("[*] Waiting 2 seconds before next call...")
        time.sleep(2)

# ============================================================================
# Summary
# ============================================================================

print("\n" + "="*80)
print("CAPTURE SUMMARY")
print("="*80)

print(f"\nKeys captured: {len(captured_data['keys'])}")
for i, key in enumerate(captured_data['keys']):
    print(f"  Key {i+1}: {key['algorithm']} - {key['hex']}")

print(f"\nIVs captured: {len(captured_data['ivs'])}")
for i, iv in enumerate(captured_data['ivs']):
    print(f"  IV {i+1}: {iv['hex']}")

print(f"\nCipher operations: {len(captured_data['encryptions'])}")
for i, enc in enumerate(captured_data['encryptions']):
    print(f"  Op {i+1}: {enc['algorithm']} - {enc['inputLength']} bytes -> {enc['outputLength']} bytes")

print("\n" + "="*80)

# Save to file
output_file = "captured_crypto.json"
with open(output_file, 'w', encoding='utf-8') as f:
    json.dump(captured_data, f, indent=2, ensure_ascii=False)

print(f"\n[+] Data saved to: {output_file}")
print("\n[+] Detaching...")
session.detach()
print("[OK] Done!\n")

